# 训练逻辑模块
# 基于Hugging Face Trainer的封装
from transformers import Trainer, TrainingArguments


def train_model(model, train_dataset, eval_dataset, training_args, compute_metrics):
    """训练模型"""
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics
    )
    return trainer.train()